from process_mongo_data_to_neo4j import (
  process_data_for_movie_info,
  process_data_for_movie_keyword,
  process_data_for_movie_crew,
  # 处理实体关系
  process_data_for_relationship_movie_info,
  process_data_for_relationship_movie_keyword,
  process_data_for_relationship_movie_crew
)
from consts import (
  # 节点
  NODE_MOVIE,
  NODE_GENRE,
  NODE_KEYWORD,
  NODE_ORIGIN_LANGUAGE,
  NODE_PRODUCTION_COMPANY,
  NODE_PRODUCTION_COUNTRY,
  NODE_CREW,
  # NODE_CAST,

  # 实体关系
  RELATIONSHIP_MOVIE_TO_GENRE,
  RELATIONSHIP_MOVIE_TO_KEYWORD,
  RELATIONSHIP_MOVIE_TO_ORIGIN_LANGUAGE,
  RELATIONSHIP_MOVIE_TO_PRODUCTION_COMPANY,
  RELATIONSHIP_MOVIE_TO_PRODUCTION_COUNTRY,
  RELATIONSHIP_MOVIE_TO_CREW
)

from mongo import (
  collection_movie_info,
  collection_movie_credits,
  collection_movie_keywords,

  get_all_data
)

from neo4j import (
  create_multiple_node,
  create_multiple_relationship_one_to_more,
  create_multiple_relationship,
  find_node_by_attr,

)

from utils import (
  record_run_time,
  log_message,
)
from log_progress import (
  LogProgress
)

# 从MongoDB中获取原始的数据
origin_data_movie_info = get_all_data(collection_movie_info)
origin_data_movie_credits = get_all_data(collection_movie_credits)
origin_data_movie_keywords = get_all_data(collection_movie_keywords)

# 无效的电影id
invalid_movie_id_set = set()

# 存储节点
# 电影info数据相关联的节点
def store_movie_info_node():
  data = process_data_for_movie_info(origin_data_movie_info, invalid_movie_id_set)
  # 存储电影节点
  create_multiple_node(NODE_MOVIE, data["movie_list"])
  # 存储类型节点
  create_multiple_node(NODE_GENRE, data["genre_list"])
  # 存储电影原始语言节点
  create_multiple_node(NODE_ORIGIN_LANGUAGE, data["origin_language_list"])
  # 存储电影制作公司节点
  create_multiple_node(NODE_PRODUCTION_COMPANY, data["production_company_list"])
  # 存储电影制作国家节点
  create_multiple_node(NODE_PRODUCTION_COUNTRY, data["production_country_list"])

# 关键词节点
def store_movie_keyword_node():
  data = process_data_for_movie_keyword(origin_data_movie_keywords, invalid_movie_id_set)
  create_multiple_node(NODE_KEYWORD, data)

# 演员+制作人员
def store_movie_crew_node():
  data = process_data_for_movie_crew(origin_data_movie_credits, invalid_movie_id_set)
  create_multiple_node(NODE_CREW, data)


# 存储实体关系
# 电影——电影类型
# 电影——电影制作公司
# 电影——电影制作国家
# 电影——电影制作人员
def store_relationship_movie_to_other_list(obj_list, obj_list_key, relationship_type, end_node_type, find_node_key='id'):
  count = 0
  total = len(obj_list)
  # 实时反馈进度和存储的时间
  log_progress = LogProgress(count, total)
  log_progress.start(relationship_type)
  node_list = []
  for item in obj_list:
    log_progress.set_count(log_progress.get_count() + 1)
    # 1.获取电影id和类型id列表
    movie_id = item['movie_id']
    other_value_list = item[obj_list_key]
    if len(other_value_list) == 0:
      continue
    # 2.根据电影id获取neo4j中的电影节点
    movie_node = find_node_by_attr(NODE_MOVIE, 'id', movie_id)
    if movie_node is None:
      continue
    # 3.根据类型id获取neo4j中的类型节点
    other_node_list = []
    for other_id in other_value_list:
      other_node = find_node_by_attr(end_node_type, find_node_key, other_id)
      if other_node is None:
        continue
      other_node_list.append(other_node)
    node_list.append({
      "start_node": movie_node,
      "end_node_list": other_node_list,
      "relationship_attrs": {}
    })
    # 4.提取节点构成列表，每300个节点存储一次
    if (log_progress.get_count() + 1) % 300 == 0:
      create_multiple_relationship_one_to_more(relationship_type, node_list)
      node_list = []
  if len(node_list):
    create_multiple_relationship_one_to_more(relationship_type, node_list)
  log_progress.close_thread()

# 电影——原始语言
def store_relationship_movie_to_origin_language(obj_list):
  count = 0
  total = len(obj_list)
  log_progress = LogProgress(count, total)
  log_progress.start(RELATIONSHIP_MOVIE_TO_ORIGIN_LANGUAGE)
  node_list = []
  for item in obj_list:
    log_progress.set_count(log_progress.get_count() + 1)
    movie_id = item['movie_id']
    original_language = item['original_language']
    movie_node = find_node_by_attr(NODE_MOVIE, 'id', movie_id)
    original_language_node = find_node_by_attr(NODE_ORIGIN_LANGUAGE, 'name', original_language)
    if movie_node is None or original_language_node is None:
      continue
    node_list.append({
      "start_node": movie_node,
      "end_node": original_language_node,
      "relationship_attrs": {}
    })
    if (log_progress.get_count() + 1) % 300 == 0:
      create_multiple_relationship(RELATIONSHIP_MOVIE_TO_ORIGIN_LANGUAGE, node_list)
      node_list = []
    if len(node_list):
      create_multiple_relationship(RELATIONSHIP_MOVIE_TO_ORIGIN_LANGUAGE, node_list)
  
  log_progress.close_thread()

# info数据与各类节点的联系
# 1.电影与类型
# 2.电影与原始语言
# 3.电影与制作公司
# 4.电影与制作国家
def store_data_for_relationship_movie_info():
  data = process_data_for_relationship_movie_info(origin_data_movie_info)
  # 存储电影与类型的实体关系
  # arg = {
  #   'obj_list': data['movie_id_to_genre_id_list'], 
  #   'obj_list_key': 'genres',
  #   'relationship_type': RELATIONSHIP_MOVIE_TO_GENRE,
  #   'end_node_type': NODE_GENRE,
  #   'find_node_key': 'id'
  # }
  # record_run_time(store_relationship_movie_to_other_list, arg)

  # 存储电影与原始语言的关系
  # arg = {'obj_list': data['movie_id_to_original_language_list']}
  # record_run_time(store_relationship_movie_to_origin_language, arg)


  # 存储电影与制作公司的实体关系
  # arg = {
  #   'obj_list': data['movie_id_to_production_company_id_list'], 
  #   'obj_list_key': 'production_companies',
  #   'relationship_type': RELATIONSHIP_MOVIE_TO_PRODUCTION_COMPANY,
  #   'end_node_type': NODE_PRODUCTION_COMPANY,
  #   'find_node_key': 'id'
  # }
  # record_run_time(store_relationship_movie_to_other_list, arg)

  # 存储电影与制作国家的实体关系
  arg = {
    'obj_list': data['movie_id_to_production_country_id_list'], 
    'obj_list_key': 'production_countries',
    'relationship_type': RELATIONSHIP_MOVIE_TO_PRODUCTION_COUNTRY,
    'end_node_type': NODE_PRODUCTION_COUNTRY,
    'find_node_key': 'iso_3166_1'
  }
  record_run_time(store_relationship_movie_to_other_list, arg)

# 存储电影与关键词
def store_data_for_relationship_movie_keyword():
  # 生成无效的id_set，以便于后续剔除不必要的数据
  # 测试时使用
  process_data_for_movie_info(origin_data_movie_info, invalid_movie_id_set)

  data = process_data_for_relationship_movie_keyword(origin_data_movie_keywords, invalid_movie_id_set)
  arg = {
    'obj_list': data, 
    'obj_list_key': 'keywords',
    'relationship_type': RELATIONSHIP_MOVIE_TO_KEYWORD,
    'end_node_type': NODE_KEYWORD,
    'find_node_key': 'id'
  }
  record_run_time(store_relationship_movie_to_other_list, arg)
  # print(len(data))
  # print(data[0])


def store_data_for_relationship_movie_crew():
  # 测试时使用
  process_data_for_movie_info(origin_data_movie_info, invalid_movie_id_set)

  # 1.存储crew关系
  data = process_data_for_relationship_movie_crew(origin_data_movie_credits, invalid_movie_id_set)
  # print(data[:3])
  # print(len(data))
  arg = {
    'obj_list': data, 
    'obj_list_key': 'crew_id_list',
    'relationship_type': RELATIONSHIP_MOVIE_TO_CREW,
    'end_node_type': NODE_CREW,
    'find_node_key': 'id'
  }
  record_run_time(store_relationship_movie_to_other_list, arg)



